%%%%%%%%%%%%
% TF analysis using Morlet wavelets
% This builds off of the pre-processing script earlier in the class
% The only change that is necessary is to make sure that the epochs are sufficiently long to resolve the frequencies that are under investigation (often longer than ERP epochs)
%%%%%%%%%%%%


%List all of the subjects

%% Calculate TF using Morlet wavelets for all subjects
subjectList={'1',
'2',
'3',
'4',
'5',
'6',
'7',
'8',
'9',
'10',
'11',
'12',
'13',
'14',
'15',
'16',
'17',
'18',
'19',
'20',
'21',
'22',
'23',
'24',
'25',
'26',
'27',
'28',
'29',
'30',
'31',
'32',
'33',
'34',
'35',
'36'}

% run a loop

openPath                                = '/Users/jsprouse/Desktop/preprocessed/'
savePath                                = '/Users/jsprouse/Desktop/N400 TF Morlet/'

for i = 1:length(subjectList)

	%clear whatever happened in the last loop, keep the subjectList and the variable i
	clearvars -except subjectList i openPath savePath
	
	%choose a subject from the list
	subject                             = char(subjectList(i));
    
    %load the file
    load([openPath subject '.N400_epoch.mat'])

    %split triggers into two conditions
    congruent                           = [55 66 77 88];
    incongruent                         = [155 166 177 188];
    
    cfg                                 = [];
    cfg.trials                          = ismember(N400_clean.trialinfo, congruent);
    data_congruent                      = ft_selectdata(cfg, N400_clean)

    cfg.trials                          = ismember(N400_clean.trialinfo, incongruent);
    data_incongruent                    = ft_selectdata(cfg, N400_clean)
	
    %pre-processing steps that are specific to TF (e.g., no baselining)
    cfg                                 = [];
    cfg.demean                          = 'yes'; %if not window, this removes the mean from the entire epoch
    
    preprocessed_congruent              = ft_preprocessing(cfg, data_congruent)
    preprocessed_incongruent            = ft_preprocessing(cfg, data_incongruent)
    
	%TF for data_unergative and data_unaccusative
	cfg = [];
	cfg.output                          = 'pow';	
	cfg.channel                         = 'all';    
	cfg.trials                          = 'all';
	cfg.method                          = 'wavelet';                
    cfg.foi                             = 2:2:40;
    cfg.toi                             = -1:0.01:1;
    cfg.width                           = 4;
    
	tf_congruent                        = ft_freqanalysis(cfg, preprocessed_congruent);
    tf_incongruent                      = ft_freqanalysis(cfg, preprocessed_incongruent);
    
	%save the results
	save([savePath subject '.N400_tf_morlet_congruent.mat'],'tf_congruent');
    save([savePath subject '.N400_tf_morlet_incongruent.mat'],'tf_incongruent');

end	

%% Calculate a grandaveraged ERP - here we eliminate a few participants with too many artifacts
subjectList={'1',
'2',
'3',
'4',
'5',
'6',
'7',
%'8',
%'9',
'10',
'11',
%'12',
'13',
'14',
%'15',
'16',
'17',
%'18',
'19',
'20',
'21',
'22',
'23',
'24',
%'25',
'26',
'27',
'28',
'29',
'30',
'31',
%'32',
'33',
'34',
'35',
'36'}

loadPath                                = '/Users/jsprouse/Desktop/N400 TF Morlet/'
savePath                                = '/Users/jsprouse/Desktop/N400 TF Morlet/'

%create an array to save the individual subject TFs
tfarray                                 = cell(2,length(subjectList));

%loop through the subjects to load their TFs
for k=1:length(subjectList)
    
    %clear whatever happened in the last loop, keep the subjectList and the variable i
	clearvars -except subjectList k loadPath savePath tfarray 
    
    %choose a subject from the list
	subject                             = char(subjectList(k));
    
    %load the files
    load([loadPath subject '.N400_tf_morlet_congruent.mat']) 
    load([loadPath subject '.N400_tf_morlet_incongruent.mat']) 
    
    %save to array
    tfarray{1,k}                        = tf_congruent;
    tfarray{2,k}                        = tf_incongruent;
    
end

%save the tfarray. They are useful sometimes (like some plots)
save([savePath 'N400_tf_morlet_array_29.mat'], 'tfarray'); 

%calculate grandaverage without keeping original (this is good for some plotting and for ft_math)
cfg                                     = [];  
cfg.keepindividual                      = 'no';  

tf_grandaverage_keepNO_congruent        = ft_freqgrandaverage(cfg, tfarray{1,:});
tf_grandaverage_keepNO_incongruent      = ft_freqgrandaverage(cfg, tfarray{2,:});

%calculate grandaverage keeping original (this is good for plotting and
%stats)
cfg                                     = [];  
cfg.keepindividual                      = 'yes';  

tf_grandaverage_keepYES_congruent       = ft_freqgrandaverage(cfg, tfarray{1,:});
tf_grandaverage_keepYES_incongruent     = ft_freqgrandaverage(cfg, tfarray{2,:});

%save both sets of grandaverages
save ([savePath 'N400_tf_morlet_grandaverage_congruent_29.mat'], 'tf_grandaverage_keepNO_congruent', 'tf_grandaverage_keepYES_congruent'); 
save ([savePath 'N400_tf_morlet_grandaverage_incongruent_29.mat'], 'tf_grandaverage_keepNO_incongruent', 'tf_grandaverage_keepYES_incongruent'); 


%% load data for plotting and statistics
loadPath                                = '/Users/jsprouse/Desktop/N400 TF Morlet/'

load([loadPath 'N400_tf_morlet_grandaverage_congruent_29.mat'])
load([loadPath 'N400_tf_morlet_grandaverage_incongruent_29.mat'])
load([loadPath 'N400_tf_morlet_array_29.mat'])

%% plot two conditions separately, all channels
%this works withe either keepNO or keepYES (the baseline function is smart,
%and will work with individuals)

cfg                                     = [];
cfg.baseline                            = [-.2 0]; 
cfg.baselinetype                        = 'db'; 

tf_grandaverage_keepNO_congruent        = ft_freqbaseline(cfg, tf_grandaverage_keepNO_congruent);
tf_grandaverage_keepNO_incongruent      = ft_freqbaseline(cfg, tf_grandaverage_keepNO_incongruent);

figure
cfg = [];
cfg.showlabels                          = 'yes'; 
cfg.fontsize                            = 14; 
cfg.layout                              = '/Users/jsprouse/Documents/MATLAB/eeglab14_1_2b/external/fieldtrip-20181210/template/layout/acticap-64ch-standard2.mat';
lay                                     = ft_prepare_layout(cfg);
cfg.xlim                                = [-0.2 1]; 
legend('congruent')
ft_multiplotTFR(cfg, tf_grandaverage_keepNO_congruent);

figure
cfg = [];
cfg.showlabels                          = 'yes'; 
cfg.fontsize                            = 14; 
cfg.layout                              = '/Users/jsprouse/Documents/MATLAB/eeglab14_1_2b/external/fieldtrip-20181210/template/layout/acticap-64ch-standard2.mat';
lay                                     = ft_prepare_layout(cfg);
cfg.xlim                                = [-0.2 1]; 
legend('incongruent')
ft_multiplotTFR(cfg, tf_grandaverage_keepNO_incongruent);

%% plot difference between two conditions, all channels
%this has to be keepNO... ft_math can't work on keepYES

%baseline
%cfg                                     = [];
%cfg.baseline                            = [-.2 0]; 
%cfg.baselinetype                        = 'db'; 

%tf_grandaverage_keepNO_congruent        = ft_freqbaseline(cfg, tf_grandaverage_keepNO_congruent);
%tf_grandaverage_keepNO_incongruent      = ft_freqbaseline(cfg, tf_grandaverage_keepNO_incongruent);

%calculate difference
cfg                                     = [];
cfg.operation                           = 'subtract';
cfg.parameter                           = 'powspctrm';
cfg.colorbar                            = 'yes';
tf_difference                           = ft_math(cfg, tf_grandaverage_keepNO_incongruent, tf_grandaverage_keepNO_congruent)

figure
cfg = [];
cfg.showlabels                          = 'yes'; 
cfg.fontsize                            = 14; 
cfg.layout                              = '/Users/jsprouse/Documents/MATLAB/eeglab14_1_2b/external/fieldtrip-20181210/template/layout/acticap-64ch-standard2.mat';
lay                                     = ft_prepare_layout(cfg);
cfg.xlim                                = [-0.2 1]; 
legend('incongruent-congruent')
ft_multiplotTFR(cfg, tf_difference);

%% plot each condition and difference between two conditions, FZ, CZ, PZ
%this has to be keepNO because of the ft_math function

%baseline
%cfg                                     = [];
%cfg.baseline                            = [-.2 0]; 
%cfg.baselinetype                        = 'db'; 

%tf_grandaverage_keepNO_congruent        = ft_freqbaseline(cfg, tf_grandaverage_keepNO_congruent);
%tf_grandaverage_keepNO_incongruent      = ft_freqbaseline(cfg, tf_grandaverage_keepNO_incongruent);

%calculate difference
cfg                                     = [];
cfg.operation                           = 'subtract';
cfg.parameter                           = 'powspctrm';
tf_difference                           = ft_math(cfg, tf_grandaverage_keepNO_incongruent, tf_grandaverage_keepNO_congruent)

%plot the figure
figure
cfg                                     = [];
cfg.parameter                           = 'powspctrm';
cfg.colorbar                            = 'yes';
cfg.showlabels                          = 'yes'; 
cfg.fontsize                            = 14; 
cfg.layout                              = '/Users/jsprouse/Documents/MATLAB/eeglab14_1_2b/external/fieldtrip-20181210/template/layout/acticap-64ch-standard2.mat';
lay                                     = ft_prepare_layout(cfg);
cfg.xlim                                = [-0.2 1];
cfg.zlim                                = [-2 2];


%plot the panels
cfg.channel                             = 'Fz'
subplot(3,3,1); ft_singleplotTFR(cfg, tf_grandaverage_keepNO_congruent);
subplot(3,3,2); ft_singleplotTFR(cfg, tf_grandaverage_keepNO_incongruent);
subplot(3,3,3); ft_singleplotTFR(cfg, tf_difference);

cfg.channel                             = 'Cz'
subplot(3,3,4); ft_singleplotTFR(cfg, tf_grandaverage_keepNO_congruent);
subplot(3,3,5); ft_singleplotTFR(cfg, tf_grandaverage_keepNO_incongruent);
subplot(3,3,6); ft_singleplotTFR(cfg, tf_difference);

cfg.channel                             = 'Pz'
subplot(3,3,7); ft_singleplotTFR(cfg, tf_grandaverage_keepNO_congruent);
subplot(3,3,8); ft_singleplotTFR(cfg, tf_grandaverage_keepNO_incongruent);
subplot(3,3,9); ft_singleplotTFR(cfg, tf_difference);

%% plotting topoplots

cfg = [];
cfg.xlim                                = [-0:0.25:.75]; %this plots a series of averages over time, where the middle number is the window size
cfg.ylim                                = [4 6];     %this averages over the frequencies                 
cfg.zlim                                = [-1.5 1.5];            
%cfg.baseline                           = [-0.2 0];          
%cfg.baselinetype                       = 'db';
cfg.layout                              = '/Users/jsprouse/Documents/MATLAB/eeglab14_1_2b/external/fieldtrip-20181210/template/layout/acticap-64ch-standard2.mat';
cfg.comment                             = 'xlim';
cfg.commentpos                          = 'title';
figure; ft_topoplotTFR(cfg,tf_difference);

%change to color brewer?
ft_hastoolbox('brewermap', 1);         % ensure this toolbox is on the path
colormap(flipud(brewermap(64,'RdBu'))) % change the colormap

%% Statistics - cluster based permutation test
%this has to be keepYES

%baseline the data
cfg = [];
cfg.baseline                            = [-.2 0]; 
cfg.baselinetype                        = 'db'; 

tf_grandaverage_keepYES_congruent       = ft_freqbaseline(cfg, tf_grandaverage_keepYES_congruent);
tf_grandaverage_keepYES_incongruent     = ft_freqbaseline(cfg, tf_grandaverage_keepYES_incongruent);

%define neighbors
cfg                                     = [];
cfg.channel                             = {'all' '-M1' '-M2'};
cfg_neighb.method                       = 'distance'; 
%cfg_neighb.neighbourdist               = 0.3;
cfg_neighb.layout                       = '/Users/jsprouse/Documents/MATLAB/eeglab14_1_2b/external/fieldtrip-20181210/template/layout/acticap-64ch-standard2.mat';
cfg_neighb.channel                      = {'all' '-M1' '-M2'};
cfg_neighb.enableedit                   = 'no';
cfg_neighb.feedback                     = 'no';	%this shows you the channel map, and then you have to click X to close it to move on.

cfg.neighbours                          = ft_prepare_neighbours(cfg_neighb);

%define design matrix for the experiment
subj                                    = size(tf_grandaverage_keepYES_incongruent.powspctrm,1);
design                                  = zeros(2,2*subj);

for i = 1:subj
  design(1,i) = i;
end

for i = 1:subj
  design(1,subj+i) = i;
end

design(2,1:subj)            = 1;
design(2,subj+1:2*subj)     = 2;

%configuration for the test
cfg.design                              = design;
cfg.uvar                                = 1; % subject number (unit variable) on line 1 of the design matrix
cfg.ivar                                = 2; %=2; condition number (independent variable) on line 2 of the design matrix

cfg.channel                             = {'all' '-M1' '-M2'};
cfg.avgoverchan                         = 'no';
cfg.avgoverfreq                         = 'no'; %if set to yes, we have to specify the frequencies to average over with cfg.frequency

cfg.method                              = 'montecarlo';
cfg.statistic                           = 'depsamplesT';
cfg.correctm                            = 'cluster';
cfg.clusteralpha                        = 0.05;
cfg.clusterstatistic                    = 'maxsum';
cfg.minnbchan                           = 2;
cfg.tail                                = 0; %two sided
cfg.clustertail                         = 0;
cfg.alpha                               = 0.025; %alpha level per tail
cfg.numrandomization                    = 500;

%tests on specific time windows
cfg.latency                             = [0 1];
full                                    = ft_freqstatistics(cfg, tf_grandaverage_keepYES_incongruent, tf_grandaverage_keepYES_congruent)

cfg.latency                             = [.3 5];
window                                  = ft_freqstatistics(cfg, tf_grandaverage_keepYES_incongruent, tf_grandaverage_keepYES_congruent)


%% plotting the statistical results? - MULTIPLOT

%multiplot
%step 1: calcluate the power difference for plotting from the keepNO
%This is the data we will plot
cfg                                     = [];
cfg.operation                           = 'subtract';
cfg.parameter                           = 'powspctrm';
tf_difference                           = ft_math(cfg, tf_grandaverage_keepNO_incongruent, tf_grandaverage_keepNO_congruent)

%specify the range of channels and the latency for the statistical test so
%that the mask from the test matches the plotting data
cfg                                     = [];
cfg.channel                             = {'all' '-M1' '-M2'};
cfg.latency                             = [0 1]

tf_difference                           = ft_selectdata(cfg, tf_difference)

%add the mask from the statistical test to the plotting data
tf_difference.mask                      = full.mask

%call the multiplot with mask settings
figure
cfg                                     = [];
cfg.channel                             = {'all' '-M1' '-M2'};
cfg.showlabels                          = 'yes'; 
cfg.fontsize                            = 14; 
cfg.layout                              = '/Users/jsprouse/Documents/MATLAB/eeglab14_1_2b/external/fieldtrip-20181210/template/layout/acticap-64ch-standard2.mat';
lay                                     = ft_prepare_layout(cfg);
cfg.xlim                                = [0 1]; 
cfg.parameter                           = 'powspctrm';
cfg.colorbar                            = 'yes';
cfg.maskparameter                       = 'mask';
cfg.maskstyle                           = 'outline';
legend('incongruent-congruent')

ft_multiplotTFR(cfg, tf_difference);


%% plotting the statistical results - CLUSTERPLOT
%clusterplot creates a topoplot... so we need to reduce dimensions
%probably want to average over frequencies
%this plots every timepoint (10ms because of the TF analysis above), so if there are a lot, it can take a long time

%average over frequencies
cfg                                     = [];
cfg.frequency                           = [4 6];
cfg.avgoverfreq                         = 'yes';

full_theta                              = ft_selectdata(cfg, full)

%make the plot
cfg                                     = [];
cfg.alpha                               = .025;            
cfg.layout                              = '/Users/jsprouse/Documents/MATLAB/eeglab14_1_2b/external/fieldtrip-20181210/template/layout/acticap-64ch-standard2.mat';
figure; ft_clusterplot(cfg, full_theta);

%change to color brewer?
ft_hastoolbox('brewermap', 1);         % ensure this toolbox is on the path
colormap(flipud(brewermap(64,'RdBu'))) 

